% 清除工作区和命令窗口
clear all;
close all;
clc;

% 参数设置
numSamples = 10000; % 样本数量
numAntennas = 4; % 天线数量
SNRdB = 20; % 信噪比 (dB)
numSubcarriers = 64; % 子载波数量
trainRatio = 0.8; % 训练集比例

% 生成通信数据
H = generateChannelData(numSamples, numAntennas, numSubcarriers);
Y = generateReceivedSignals(H, SNRdB);

% 划分训练集和测试集
trainSize = floor(trainRatio * numSamples);
trainH = H(1:trainSize, :, :);
trainY = Y(1:trainSize, :, :);
testH = H(trainSize+1:end, :, :);
testY = Y(trainSize+1:end, :, :);

% 数据预处理
trainY = reshape(trainY, [trainSize, numAntennas * numSubcarriers]);
testY = reshape(testY, [numSamples - trainSize, numAntennas * numSubcarriers]);
trainH = reshape(trainH, [trainSize, numAntennas * numSubcarriers]);
testH = reshape(testH, [numSamples - trainSize, numAntennas * numSubcarriers]);

% 归一化数据
trainY = normalize(trainY);
testY = normalize(testY);
trainH = normalize(trainH);
testH = normalize(testH);

% 构建深度学习模型
layers = [
    featureInputLayer(numAntennas * numSubcarriers, 'Name', 'input')
    fullyConnectedLayer(256, 'Name', 'fc1')
    reluLayer('Name', 'relu1')
    fullyConnectedLayer(128, 'Name', 'fc2')
    reluLayer('Name', 'relu2')
    fullyConnectedLayer(numAntennas * numSubcarriers, 'Name', 'output')
    regressionLayer('Name', 'regression')
];

options = trainingOptions('adam', ...
    'MaxEpochs', 50, ...
    'MiniBatchSize', 64, ...
    'Shuffle', 'every-epoch', ...
    'Verbose', false, ...
    'Plots', 'training-progress');

net = trainNetwork(trainY, trainH', layers, options);

% 进行信道估计
estimatedH = predict(net, testY');
estimatedH = reshape(estimatedH', [numSamples - trainSize, numAntennas, numSubcarriers]);
testH = reshape(testH, [numSamples - trainSize, numAntennas, numSubcarriers]);

% 性能评估
% 均方误差 (MSE)
mse = mean(mean(mean((abs(estimatedH - testH)).^2, 3), 2), 1);
fprintf('均方误差 (MSE): %.6f\n', mse);

% 均方根误差 (RMSE)
rmse = sqrt(mse);
fprintf('均方根误差 (RMSE): %.6f\n', rmse);

% 峰值信噪比 (PSNR)
maxVal = max(max(max(abs(testH))));
psnr = 20 * log10(maxVal / rmse);
fprintf('峰值信噪比 (PSNR): %.6f dB\n', psnr);

% 结构相似性指数 (SSIM)
ssimVals = zeros(size(testH, 1), 1);
for i = 1:size(testH, 1)
    ssimVals(i) = ssim2(squeeze(abs(testH(i, :, :))), squeeze(abs(estimatedH(i, :, :))));
end
meanSSIM = mean(ssimVals);
fprintf('平均结构相似性指数 (SSIM): %.6f\n', meanSSIM);

% 绘制真实信道和估计信道的对比图
figure;
subplot(2,1,1);
imagesc(abs(testH(1,:,:)));
title('真实信道');
colorbar;

subplot(2,1,2);
imagesc(abs(estimatedH(1,:,:)));
title('估计信道');
colorbar;

% 生成信道数据的函数
function H = generateChannelData(numSamples, numAntennas, numSubcarriers)
    H = zeros(numSamples, numAntennas, numSubcarriers);
    for i = 1:numSamples
        for j = 1:numAntennas
            for k = 1:numSubcarriers
                H(i, j, k) = (randn + 1j*randn)/sqrt(2);
            end
        end
    end
end

% 生成接收信号的函数
function Y = generateReceivedSignals(H, SNRdB)
    numSamples = size(H, 1);
    numAntennas = size(H, 2);
    numSubcarriers = size(H, 3);
    Y = zeros(numSamples, numAntennas, numSubcarriers);
    for i = 1:numSamples
        for j = 1:numAntennas
            for k = 1:numSubcarriers
                s = (randn + 1j*randn)/sqrt(2); % 发送信号
                noisePower = 10^(-SNRdB/10);
                noise = (randn + 1j*randn)*sqrt(noisePower/2);
                Y(i, j, k) = H(i, j, k) * s + noise;
            end
        end
    end
end

% 归一化数据的函数
function normalizedData = normalize(data)
    meanData = mean(data, 1);
    stdData = std(data, 0, 1);
    normalizedData = (data - repmat(meanData, size(data, 1), 1))./repmat(stdData, size(data, 1), 1);
end
    